#include <stdio.h>
#include <stdint.h>
#include <stdlib.h>
#include <time.h>
#include <windows.h>

#define FILL_FILE_NAME "%s\\_DISKFILL_%u.DAT"
#define FILL_FILE_INFO "%s\\_DISKFILL_*.DAT"
#define TEMP_FILE_SIZE (1UL * 1024 * 1024 * 1024)
#define DATA_MAGIC 0x1234abcd5678dcbaULL

struct {
    char disk[16];
    uint64_t total_size;
    uint64_t aval_size;
    uint64_t free_size;
    uint64_t actual_size;
    enum {
        STATUS_IDLE,
        STATUS_INITING,
        STATUS_FILLING,
        STATUS_CHECKING,
        STATUS_DONE,
        STATUS_ERROR
    } status;
} g_args = {0};

void print_usage(const char *cmd_name)
{
    if(cmd_name == NULL)
        cmd_name = "DiskFill";
    printf("Usage:\n"
           "    %s <disk label>\n"
           "Example:\n"
           "    %s C:\n\n",
           cmd_name, cmd_name);
}

const char* parse_args(int argc, char **argv)
{
    int i;

    if(argc != 2)
    {
        print_usage(argv[0]);
        return "Missing argument <disk label>";
    }

    for(i = 0; i < 8 && argv[1][i] != '\0'; i++)
        g_args.disk[i] = argv[1][i];
    g_args.disk[i] = '\0';

    if(i < 2 || g_args.disk[i - 1] != ':')
    {
        print_usage(argv[0]);
        return "Argument 1 is not disk label.";
    }

    return NULL;
}

int format_size(uint64_t size, char *out, int out_len)
{
    int len;

    if(size < 1000) // B
    {
        len = sprintf_s(out, out_len, "%uB", size);
    }
    else if(size < 1000 * 1024) // KB
    {
        len = sprintf_s(out, out_len, "%.3fKB", (float)size / 1024);
    }
    else if(size < 1000 * 1024 * 1024) // MB
    {
        len = sprintf_s(out, out_len, "%.3fMB", (float)size / (1024 * 1024));
    }
    else if(size < ((uint64_t)1000) * 1024 * 1024 * 1024) // GB
    {
        len = sprintf_s(out, out_len, "%.3fGB", (float)size / (1024 * 1024 * 1024));
    }
    else // TB
    {
        len = sprintf_s(out, out_len, "%.3fTB", (float)size / (1024ull * 1024 * 1024 * 1024));
    }

    if(len <= 0)
    {
        out[0] = '\0';
        printf("Error: format size failed. ret=%d\n", len);
        return -1;
    }
    return 0;
}

int format_time(uint64_t second, char *out, int out_len)
{
    int ret;

    if(second < 60)
        ret = sprintf_s(out, out_len, "0:%02u", second);
    else if(second < 60 * 60)
        ret = sprintf_s(out, out_len, "%u:%02u", second / 60, second % 60);
    else
        ret = sprintf_s(out, out_len, "%u:%02u:%02u", second / (60 * 60), (second / 60) % 60, second % 60);

    if(ret <= 0)
    {
        out[0] = '\0';
        printf("Error: format time failed. ret=%d\n", ret);
        return -1;
    }
    return 0;
}

int get_disk_info()
{
    DWORD64 aval_byte = 0;
    DWORD64 total_byte = 0;
    DWORD64 free_byte = 0;
    BOOL ret;

    ret = GetDiskFreeSpaceEx((LPSTR)g_args.disk, (PULARGE_INTEGER)&aval_byte, (PULARGE_INTEGER)&total_byte, (PULARGE_INTEGER)&free_byte);
    if(ret != TRUE)
    {
        printf("Error: get disk information failed. disk='%s'\n", g_args.disk);
        return -1;
    }
    if(total_byte == 0)
    {
        printf("Error: Disk '%s' is not a storage device.\n", g_args.disk);
        return -1;
    }

    g_args.aval_size = aval_byte;
    g_args.total_size = total_byte;
    g_args.free_size = free_byte;

    return 0;
}

uint8_t *get_memblock(uint64_t size)
{
    uint64_t *mem = NULL;
    uint32_t block_num;
    uint32_t i;

    block_num = (size + sizeof(*mem) - 1) / sizeof(*mem);
    size = block_num * sizeof(*mem);
    mem = malloc(size);
    if(mem == NULL)
    {
        printf("Error: allocate memory block failed. size=%u\n", size);
        return NULL;
    }

    for(i = 0; i < block_num; i++)
    {
        mem[i] = DATA_MAGIC;
    }

    return (void*)mem;
}

int get_file_name(const char *disk, uint32_t file_index, char *out, uint32_t out_len)
{
    int len = sprintf_s(out, out_len, FILL_FILE_NAME, disk, file_index);
    if(len <= 0)
    {
        printf("Error: format file name failed. disk='%s' index=%u\n", disk, file_index);
        return -1;
    }

    return 0;
}

volatile struct {
    uint64_t filled_size;
    uint32_t file_num;
    uint32_t file_total;
    uint64_t checked_size;
} g_work_progess = {0};

int fill_file(const char *file_path, uint8_t *mem, uint64_t mem_size, uint64_t file_max_size)
{
    FILE *file = NULL;
    uint64_t filled_size = 0;
    uint64_t write_size = 0;
    uint64_t ret;

    file = fopen(file_path, "wb");
    if(file == NULL)
    {
        printf("Error: can not create file. file='%s'\n", file_path);
        return -1;
    }

    // file created, add count
    g_work_progess.file_num ++;
    while(filled_size < file_max_size)
    {
        if(filled_size + mem_size > file_max_size)
            write_size = file_max_size - filled_size;
        else
            write_size = mem_size;

        ret = fwrite(mem, 1, write_size, file);
        if(ret > 0)
        {
            g_work_progess.filled_size += ret;
            filled_size += ret;
        }

        if(ret < mem_size) // disk error or disk full
        {
            fclose(file);
            return 1;
        }
    }

    fclose(file);
    return 0;
}

DWORD fill_disk(LPVOID param)
{
    char file_path[64] = {0};
    uint8_t *mem = NULL;
    uint64_t file_size = 0;
    uint64_t file_max_size = TEMP_FILE_SIZE;
    const uint64_t block_size = 1ull * 1024 * 1024; // 1MB
    uint32_t i;
    int ret;

    if(g_args.free_size == 0)
    {
        g_args.status = STATUS_DONE;
        return 0;
    }

    mem = get_memblock(block_size);
    if(mem == NULL)
    {
        g_args.status = STATUS_ERROR;
        return 0;
    }

    g_work_progess.filled_size = 0;
    g_work_progess.file_total = (g_args.free_size + file_max_size - 1) / file_max_size;
    g_work_progess.file_num = 0;

    for(i = 0; i < g_work_progess.file_total; i++)
    {
        ret =get_file_name(g_args.disk, i, file_path, sizeof(file_path));
        if(ret != 0)
        {
            free(mem);
            g_args.status = STATUS_ERROR;
            return 0;
        }

        if(g_work_progess.filled_size + file_max_size > g_args.free_size)
            file_max_size = g_args.free_size - g_work_progess.filled_size;
        else
            file_max_size = TEMP_FILE_SIZE;
        ret = fill_file(file_path, mem, block_size, file_max_size);
        if(ret == -1)
        {
            free(mem);
            g_args.status = STATUS_ERROR;
            return 0;
        }
        if(ret == 1)
            break;
    }

    g_args.status = STATUS_DONE;
    free(mem);
    return 0;
}

void show_fill_progress()
{
    uint64_t action_start_time = clock();
    uint64_t start_time = action_start_time;
    uint64_t start_size = g_work_progess.filled_size;
    uint64_t end_time;
    uint64_t end_size;

    char total_str[32] = {0};
    char filled_str[32] = {0};
    char rate_str[32] = {0};
    char time_str[32] = {0};
    uint64_t rate;
    float progress;

    if(g_args.free_size == 0)
        return;

    (void)format_size(g_args.free_size, total_str, sizeof(total_str));

    printf("Filling...");
    while(g_args.status == STATUS_FILLING)
    {
        Sleep(1000);
        end_time = clock();
        end_size = g_work_progess.filled_size;
        if(end_time == start_time)
        {
            end_time++;
        }

        if(end_size == 0)
            progress = 0;
        else
            progress = (float)(end_size * 100) / g_args.free_size;

        rate = (end_size - start_size) * CLOCKS_PER_SEC / (end_time - start_time);

        (void)format_size(end_size, filled_str, sizeof(filled_str));
        (void)format_size(rate, rate_str, sizeof(rate_str));
        (void)format_time((end_time - action_start_time) / CLOCKS_PER_SEC, time_str, sizeof(time_str));
        printf("\rProgress: %10s/%-10s  %4.1f%%  Rate:%10s/s Passed: %10s",
                filled_str, total_str, progress, rate_str, time_str);

        start_time = clock();
        start_size = g_work_progess.filled_size;
    }
    printf("\nFilling disk finished.\n");
}

int do_fill()
{
    DWORD threadId = 0;
    HANDLE handle;

    g_args.status = STATUS_FILLING;
    memset((void*)&g_work_progess, 0, sizeof(g_work_progess));
    handle = CreateThread(NULL, 0, fill_disk, NULL, 0, &threadId);
    if(handle == NULL)
    {
        printf("Error: create thread failed.");
        return -1;
    }
    Sleep(500);
    show_fill_progress();
    WaitForSingleObject(handle, INFINITE);

    return 0;
}

int check_file(const char *file_path, uint8_t *mem, uint64_t mem_size)
{
    FILE *file = NULL;
    int ret;
    long int align_len = 0;
    uint32_t i;

    file = fopen(file_path, "rb");
    if(file == NULL)
    {
        printf("Error: open file failed. file='%s'\n", file_path);
        return -1;
    }

    while(feof(file) == 0)
    {
        ret = fread(mem, 1, mem_size, file);
        if(ret < 0)
        {
            printf("Error: read file failed. file='%s' read=%d \n", file_path, ret);
            fclose(file);
            return 0;
        }

        for(i = align_len; i < (uint32_t)ret; i += sizeof(uint64_t))
        {
            if(*((uint64_t*)(mem + i)) != DATA_MAGIC)
            {
                printf("\nError: data mismatch. data=%016llx.\n", *((uint64_t*)(mem + i)));
                if(mem[i] != 0x0 && mem[i] != 0xff)
                {
                    printf("       position=%u len=%u\n", i, ret);
                }
                fclose(file);
                return 1;
            }
        }

        // align to sizeof(uint64_t)
        align_len = sizeof(uint64_t) - (ret % sizeof(uint64_t));
        g_work_progess.checked_size += ret;
    }

    fclose(file);
    return 0;
}

DWORD check_disk(LPVOID param)
{
    const uint64_t block_size = 1ull * 1024 * 1024; // 1MB
    char file_path[64] = {0};
    uint8_t *mem = NULL;
    int ret;
    uint32_t i;

    if(g_work_progess.file_num == 0)
    {
        g_args.status = STATUS_DONE;
        return 0;
    }

    mem = get_memblock(block_size);
    if(mem == NULL)
    {
        g_args.status = STATUS_ERROR;
        return 0;
    }

    for(i = 0; i < g_work_progess.file_num; i++)
    {
        ret = get_file_name(g_args.disk, i, file_path, sizeof(file_path));
        if(ret != 0)
        {
            free(mem);
            g_args.status = STATUS_ERROR;
            return 0;
        }

        ret = check_file(file_path, mem, block_size);
        if(ret != 0)
        {
            free(mem);
            g_args.status = STATUS_ERROR;
            return 0;
        }
    }

    free(mem);
    g_args.status = STATUS_DONE;
    return 0;
}

void show_check_progress()
{
    uint64_t action_start_time = clock();
    uint64_t start_time = action_start_time;
    uint64_t start_size = g_work_progess.checked_size;
    uint64_t end_time;
    uint64_t end_size;

    char total_str[32] = {0};
    char checkeded_str[32] = {0};
    char rate_str[32] = {0};
    char time_str[32] = {0};
    uint64_t rate;
    float progress;

    if(g_work_progess.filled_size == 0)
        return;

    (void)format_size(g_work_progess.filled_size, total_str, sizeof(total_str));

    printf("Checking...");
    while(g_args.status == STATUS_CHECKING)
    {
        Sleep(1000);
        end_time = clock();
        end_size = g_work_progess.checked_size;
        if(end_time == start_time)
        {
            end_time++;
        }

        if(end_size == 0)
            progress = 0;
        else
            progress = (float)(end_size * 100) / g_work_progess.filled_size;

        rate = (end_size - start_size) * CLOCKS_PER_SEC / (end_time - start_time);

        (void)format_size(end_size, checkeded_str, sizeof(checkeded_str));
        (void)format_size(rate, rate_str, sizeof(rate_str));
        (void)format_time((end_time - action_start_time) / CLOCKS_PER_SEC, time_str, sizeof(time_str));
        printf("\rProgress: %10s/%-10s  %4.1f%%  Rate:%10s/s Passed: %10s",
                checkeded_str, total_str, progress, rate_str, time_str);

        start_time = clock();
        start_size = g_work_progess.checked_size;
    }
    printf("\nChecking disk finished.\n");
}

int do_check()
{
    DWORD threadId = 0;
    HANDLE handle;

    g_args.status = STATUS_CHECKING;
    handle = CreateThread(NULL, 0, check_disk, NULL, 0, &threadId);
    if(handle == NULL)
    {
        printf("Error: create thread failed.");
        return -1;
    }
    Sleep(500);
    show_check_progress();
    WaitForSingleObject(handle, INFINITE);

    return 0;
}

int do_delete()
{
    char file_path[64] = {0};
    uint8_t *mem = NULL;
    int ret;
    uint32_t i;

    printf("Deleting...");
    for(i = 0; i < g_work_progess.file_num; i++)
    {
        printf("\rProgress: %u/%u", i + 1, g_work_progess.file_num);
        ret = get_file_name(g_args.disk, i, file_path, sizeof(file_path));
        if(ret != 0)
            continue;

        ret = DeleteFile(file_path);
        if(!ret)
        {
            printf("\nError: delete file failed.\n"
                   "       Please delete file '%s' manually later.\n", file_path);
            continue;
        }
    }

    printf("\rDeleting file finished.\n");
    return 0;
}

void show_check_report()
{
    char total_str[32] = {0};
    char usable_free_str[32] = {0};
    char usable_str[32] = {0};
    const char *result_str = NULL;

    (void)format_size(g_args.free_size, total_str, sizeof(total_str));
    (void)format_size(g_work_progess.checked_size, usable_free_str, sizeof(usable_free_str));
    (void)format_size(g_work_progess.checked_size + g_args.total_size - g_args.free_size, usable_str, sizeof(usable_str));
    if(g_work_progess.filled_size == g_args.free_size)
        result_str = "PASS";
    else
        result_str = "FAIL";


    printf("\nCheck report:\n"
           "    Total free space:  %s\n"
           "    Usable free space: %s\n"
           "    Usable space:      %s\n"
           "    Create files:      %u\n"
           "\n%s\n",
           total_str, usable_free_str, usable_str, g_work_progess.file_num, result_str);
}

int print_disk_info()
{
    char size_str[32];
    int ret;

    puts("");
    printf("Disk space of %s\n", g_args.disk);

    ret = format_size(g_args.total_size, size_str, sizeof(size_str));
    if(ret != 0)
        return -1;
    printf("Total: %10s\n", size_str);

    ret = format_size(g_args.aval_size, size_str, sizeof(size_str));
    if(ret != 0)
        return -1;
    printf("Avail: %10s\n", size_str);

    ret = format_size(g_args.free_size, size_str, sizeof(size_str));
    if(ret != 0)
        return -1;
    printf("Free:  %10s\n", size_str);

    return 0;
}

int wait_select(const char *info)
{
    int ret;
    int selected = 0;

    printf("%s (Y or N):", info);
    while(selected == 0)
    {
        ret = getchar();
        if((ret | 32) == 'y' || (ret | 32) == 'n')
        {
            selected = (ret | 32);
            ret = getchar();
            if(ret != '\n' && ret != '\r')
                selected = 0;
        }

        while(ret != '\n')
            ret = getchar();
        if(selected == 0)
            printf("Please input Y or N:");
    }

    if(selected == 'y')
        return 1;

    return 0;
}

int main(int argc, char **argv)
{
    const char *errMsg;
    int ret;

    errMsg = parse_args(argc, argv);
    if(errMsg != NULL)
    {
        puts(errMsg);
        return 1;
    }

    ret = get_disk_info();
    if(ret != 0)
        return 1;

    ret = print_disk_info();
    if(ret != 0)
        return 1;
    printf("\nChecking steps:\n"
            "    1. Fill disk with temporary files.\n"
            "    2. Check file blocks.\n"
            "    3. Delete temporary files.\n");
    ret = wait_select("Sure to fill disk?");
    if(ret == 0)
    {
        printf("Aborted.\n");
        return 0;
    }

    ret = do_fill();
    if(ret != 0 || g_args.status == STATUS_ERROR)
    {
        printf("An error occurred.\n"
                "    You can delete files " FILL_FILE_INFO " later if exists.\n"
               "Aborted.\n", g_args.disk);
        return 1;
    }

    printf("\nSuggest that plugging out the disk and plugging in again to clean the cache.\n");
    ret = wait_select("Sure to check?");
    if(ret == 0)
    {
        printf("Aborted.\n"
               "    You can delete files " FILL_FILE_INFO " later if exists.\n", g_args.disk);
        return 0;
    }

    ret = do_check();
    if(ret != 0 || g_args.status == STATUS_ERROR)
    {
        printf("An error occurred.\n    Maybe the temporary files ware modified.\n");
    }

    ret = wait_select("Sure to delete temporary files?");
    if(ret == 0)
    {
        printf("Aborted.\n"
               "    You can delete files " FILL_FILE_INFO " later if exists.\n", g_args.disk);
        show_check_report();
        return 0;
    }

    do_delete();
    show_check_report();

    return 0;
}
